package rdd

import org.apache.spark.rdd.RDD
import org.apache.spark.{HashPartitioner, Partitioner, SparkConf, SparkContext}

object RDD_PartitionsTest {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf()
    conf.setMaster("local[*]")
    conf.setAppName("RDD_PartitionsTest")

    val sc = new SparkContext(conf)
    val pairRDD: RDD[(Int, Int)] = sc.parallelize(List((1, 1), (2, 2), (3, 3)),2)
    println(pairRDD.partitions.size)

    // 对RDD进行重分区
    val rdd1 = pairRDD.repartition(4)
    println(rdd1.partitions.size)
    val rdd2 = pairRDD.coalesce(1, shuffle = true)
    println(rdd2.partitions.size)

    // partitionBy()方法重分区
    val rdd3 = pairRDD.partitionBy(new HashPartitioner(6))
    println(rdd3.partitions.size)

    // 自定义数据分区器
    val rdd10: RDD[Int] = sc.parallelize(1 to 100000)
    rdd10
      .map(number => (number,null))
      .partitionBy(new MyPartitioner(2))
      .map(_._1)
      .saveAsTextFile("output")

    sc.stop()
  }


  class MyPartitioner(numberOfPartitions:Int) extends Partitioner{
    override def numPartitions: Int = numberOfPartitions

    override def getPartition(key: Any): Int = {
      key.toString.toInt % numPartitions
    }
  }

}
